// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#include "FieldManagerBuilder.hpp"

namespace SiYuan {

//=======================================================================
namespace {
  struct PostRegistrationFunctor {

    const std::vector<bool>& active_;
    PHX::FieldManager<panzer::Traits>& fm_;
    panzer::Traits::SD& setup_data_;

    PostRegistrationFunctor(const std::vector<bool>& active,
                            PHX::FieldManager<panzer::Traits>& fm,
                            panzer::Traits::SD& setup_data)
      : active_(active),fm_(fm),setup_data_(setup_data) {}

    template<typename T>
    void operator()(T) const {
      auto index = Sacado::mpl::find<panzer::Traits::EvalTypes,T>::value;
      if (active_[index])
        fm_.postRegistrationSetupForType<T>(setup_data_);
    }
  };
}

FieldManagerBuilder::
FieldManagerBuilder(bool disablePhysicsBlockScatter,
                    bool disablePhysicsBlockGather)
  : panzer::FieldManagerBuilder(disablePhysicsBlockScatter,disablePhysicsBlockScatter)
{}

/** This class defines a response based on a functional. */
  
void FieldManagerBuilder::setupVolumeFieldManagers( const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh,
              const std::vector<Teuchos::RCP<panzer::PhysicsBlock> >& physicsBlocks,
              const std::vector<panzer::WorksetDescriptor> & wkstDesc,
					    const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& cm_factory,
					    const Teuchos::ParameterList& closure_models,
              const panzer::LinearObjFactory<panzer::Traits> & lo_factory,
					    const Teuchos::ParameterList& user_data,
              const panzer::GenericEvaluatorFactory & gEvalFact,
              bool closureModelByEBlock)
{
  TEUCHOS_TEST_FOR_EXCEPTION(getWorksetContainer()==Teuchos::null,std::logic_error,
                            "panzer::FMB::setupVolumeFieldManagers: method function getWorksetContainer() returns null. "
                            "Plase call setWorksetContainer() before calling this method");
  TEUCHOS_TEST_FOR_EXCEPTION(physicsBlocks.size()!=wkstDesc.size(),std::runtime_error,
                            "panzer::FMB::setupVolumeFieldManagers: physics block count must match workset descriptor count.");

  phx_volume_field_managers_.clear();

  Teuchos::RCP<const panzer::GlobalIndexer> globalIndexer = lo_factory.getRangeGlobalIndexer();

  for (std::size_t blkInd=0;blkInd<physicsBlocks.size();++blkInd) {
    Teuchos::RCP<PhysicsBlock> pb = Teuchos::rcp_dynamic_cast<PhysicsBlock>( physicsBlocks[blkInd] );
    const panzer::WorksetDescriptor wd = wkstDesc[blkInd];

    panzer::Traits::SD setupData;
    setupData.worksets_ = getWorksetContainer()->getWorksets(wd);
    setupData.orientations_ = getWorksetContainer()->getOrientations();
    if(setupData.worksets_->size()==0)
      continue;

    // sanity check
    TEUCHOS_ASSERT(wd.getElementBlock()==pb->elementBlockID());

    // build a field manager object
    Teuchos::RCP<PHX::FieldManager<panzer::Traits> > fm
          = Teuchos::rcp(new PHX::FieldManager<panzer::Traits>);

    // use the physics block to register active evaluators
    pb->setActiveEvaluationTypes(active_evaluation_types_);
    pb->buildAndRegisterEquationSetEvaluators(*fm, user_data);
    pb->setMaterial(*fm, _materials);
    if(!physicsBlockGatherDisabled())
      pb->buildAndRegisterGatherAndOrientationEvaluators(*fm,lo_factory,user_data);
    pb->buildAndRegisterDOFProjectionsToIPEvaluators(*fm,Teuchos::ptrFromRef(lo_factory),user_data);
    if(!physicsBlockScatterDisabled())
      pb->buildAndRegisterScatterEvaluators(*fm,lo_factory,user_data);

    if(closureModelByEBlock)
      pb->buildAndRegisterClosureModelEvaluators(*fm,cm_factory,pb->elementBlockID(),closure_models,user_data);
    else
      pb->buildAndRegisterClosureModelEvaluators(*fm,cm_factory,closure_models,user_data);
    
    // Reset active evaluation types
    pb->activateAllEvaluationTypes();

    // register additional model evaluator from the generic evaluator factory
    gEvalFact.registerEvaluators(*fm,wd,*pb);

    // setup derivative information
    setKokkosExtendedDataTypeDimensions(wd.getElementBlock(),*globalIndexer,user_data,*fm);

    // call postRegistrationSetup() for each active type
    Sacado::mpl::for_each_no_kokkos<panzer::Traits::EvalTypes>(PostRegistrationFunctor(active_evaluation_types_,*fm,setupData));

    // make sure to add the field manager & workset to the list
    volume_workset_desc_.push_back(wd);
    phx_volume_field_managers_.push_back(fm);
  }
}

//=======================================================================
void FieldManagerBuilder::setupVolumeFieldManagers( const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh,
              const std::vector<Teuchos::RCP<panzer::PhysicsBlock> >& physicsBlocks,
					    const panzer::ClosureModelFactory_TemplateManager<panzer::Traits>& cm_factory,
					    const Teuchos::ParameterList& closure_models,
              const panzer::LinearObjFactory<panzer::Traits> & lo_factory,
					    const Teuchos::ParameterList& user_data)
{
   std::vector<panzer::WorksetDescriptor> wkstDesc;
   for(std::size_t i=0;i<physicsBlocks.size();i++)
     wkstDesc.push_back(panzer::blockDescriptor(physicsBlocks[i]->elementBlockID()));

   panzer::EmptyEvaluatorFactory eef;
   setupVolumeFieldManagers(mesh,physicsBlocks,wkstDesc,cm_factory,closure_models,lo_factory,user_data,eef);
}

//=======================================================================
void FieldManagerBuilder::setupDirichletFieldManager(const std::vector< Dirichlet<panzer::Traits::RealType> >& ds,
              const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh )
{
  phx_dirichlet_field_manager_ = std::shared_ptr<PHX::FieldManager<panzer::Traits>>( new PHX::FieldManager<panzer::Traits>());
  if( ds.empty() ) return;

  Teuchos::RCP< DirichletsEvalutor<panzer::Traits::Residual, panzer::Traits> > re =
    Teuchos::rcp( new DirichletsEvalutor<panzer::Traits::Residual, panzer::Traits>(ds, mesh) );
  phx_dirichlet_field_manager_->registerEvaluator<panzer::Traits::Residual>(re);
  phx_dirichlet_field_manager_->requireField<panzer::Traits::Residual>(*re->evaluatedFields()[0]);

  Teuchos::RCP< DirichletsEvalutor<panzer::Traits::Jacobian, panzer::Traits> > je =
    Teuchos::rcp( new DirichletsEvalutor<panzer::Traits::Jacobian, panzer::Traits>(ds, mesh) );
  phx_dirichlet_field_manager_->registerEvaluator<panzer::Traits::Jacobian>(je);
  phx_dirichlet_field_manager_->requireField<panzer::Traits::Jacobian>(*je->evaluatedFields()[0]);

  Teuchos::RCP< DirichletsEvalutor<panzer::Traits::Tangent, panzer::Traits> > te =
    Teuchos::rcp( new DirichletsEvalutor<panzer::Traits::Tangent, panzer::Traits>(ds, mesh) );
  phx_dirichlet_field_manager_->registerEvaluator<panzer::Traits::Tangent>(te);
  phx_dirichlet_field_manager_->requireField<panzer::Traits::Tangent>(*te->evaluatedFields()[0]);

  panzer::Traits::SD setupData;

//  std::vector<PHX::index_size_type> derivative_dimensions;
//  derivative_dimensions.push_back(8);
//  phx_dirichlet_field_manager_->setKokkosExtendedDataTypeDimensions<panzer::Traits::Jacobian>(derivative_dimensions);
//  phx_dirichlet_field_manager_->setKokkosExtendedDataTypeDimensions<panzer::Traits::Tangent>(derivative_dimensions);
//  Sacado::mpl::for_each_no_kokkos<panzer::Traits::EvalTypes>( PostRegistrationFunctor
//      (active_evaluation_types_, *phx_dirichlet_field_manager_, setupData) );
  phx_dirichlet_field_manager_->postRegistrationSetup(setupData);
}

//=======================================================================
//=======================================================================
void FieldManagerBuilder::setupNeumannFieldManager(const std::vector< NeumannBoundary >& ns,
  const std::vector< CLoad >& cloads,
  const panzer::LinearObjFactory<panzer::Traits> & lo_factory,
  const std::vector<Teuchos::RCP<panzer::PhysicsBlock> >& physicsBlocks )
{
  TEUCHOS_TEST_FOR_EXCEPTION(getWorksetContainer()==Teuchos::null,std::logic_error,
                            "SiYuan::setupNeumannFieldManagers: method function getWorksetContainer() returns null. "
                            "Plase call setWorksetContainer() before calling this method");

  phx_neumann_field_manager_ = std::shared_ptr<PHX::FieldManager<panzer::Traits>>( new PHX::FieldManager<panzer::Traits>());
  if( ns.empty() && cloads.empty() ) return;

  const Teuchos::RCP<const panzer::GlobalIndexer> globalIndexer = lo_factory.getRangeGlobalIndexer();

  std::map<std::string,Teuchos::RCP<panzer::PhysicsBlock> > physicsBlocks_map;
  if( !ns.empty() ) {
  // for convenience build a map (element block id => physics block)
     std::vector<Teuchos::RCP<panzer::PhysicsBlock> >::const_iterator blkItr;
     for(blkItr=physicsBlocks.begin();blkItr!=physicsBlocks.end();++blkItr) {
        Teuchos::RCP<panzer::PhysicsBlock> pb = *blkItr;
        std::string blockId = pb->elementBlockID();

        // add block id, physics block pair to the map
        physicsBlocks_map.insert(std::make_pair(blockId,pb));
     }
  }

  for(const auto bc : ns) {
    panzer::WorksetDescriptor wd = bc.getWorksetDescriptor();
    const Teuchos::RCP<std::map<unsigned,panzer::Workset> >
      currentWkst = getWorksetContainer()->getSideWorksets(wd);
    if (currentWkst.is_null()) continue;

    for (std::map<unsigned,panzer::Workset>::const_iterator wkst = currentWkst->begin();
           wkst != currentWkst->end(); ++wkst)
    {
      std::map<std::string,Teuchos::RCP<panzer::PhysicsBlock> >::const_iterator
            volume_pb_itr = physicsBlocks_map.find(bc.m_eblock_name);
      TEUCHOS_TEST_FOR_EXCEPTION(volume_pb_itr == physicsBlocks_map.end(), std::logic_error,
            "panzer::FMB::setupBCFieldManagers: Cannot find physics block corresponding to element block \""
            << bc.m_eblock_name << "\"");

      const Teuchos::RCP<const panzer::PhysicsBlock> volume_pb = physicsBlocks_map.find(bc.m_eblock_name)->second;
      const Teuchos::RCP<const shards::CellTopology> volume_cell_topology = volume_pb->cellData().getCellTopology();

      // register evaluators from strategy
      const panzer::CellData side_cell_data(wkst->second.num_cells,
                wkst->second.details(0).subcell_index, volume_cell_topology);

      // Copy the physics block for side integrations
      Teuchos::RCP<panzer::PhysicsBlock> side_pb = volume_pb->copyWithCellData(side_cell_data);

      //Sacado::mpl::for_each_no_kokkos<panzer::Traits::EvalTypes>
      Teuchos::RCP< NeumannEvalutor<panzer::Traits::Residual, panzer::Traits> > re =
        Teuchos::rcp( new NeumannEvalutor<panzer::Traits::Residual, panzer::Traits>(bc,globalIndexer,side_pb) );
      phx_neumann_field_manager_->registerEvaluator<panzer::Traits::Residual>(re);
      phx_neumann_field_manager_->requireField<panzer::Traits::Residual>(*re->evaluatedFields()[0]);

    /*  Teuchos::RCP< NeumannEvalutor<panzer::Traits::Jacobian> > je =
        Teuchos::rcp( new NeumannEvalutor<panzer::Traits::Jacobian>(bc,mesh,*side_pb) );
      phx_neumann_field_manager_->registerEvaluator<panzer::Traits::Jacobian>(je);
      phx_neumann_field_manager_->requireField<panzer::Traits::Jacobian>(*je->evaluatedFields()[0]);*/
    }
  }

  if( !cloads.empty() ) {
    Teuchos::RCP< PHX::Evaluator<panzer::Traits> > re =
      Teuchos::rcp( new CLoadEvalutor<panzer::Traits::Residual, panzer::Traits>(cloads,globalIndexer) );
    phx_neumann_field_manager_->registerEvaluator<panzer::Traits::Residual>(re);
    phx_neumann_field_manager_->requireField<panzer::Traits::Residual>(*re->evaluatedFields()[0]);

    Teuchos::RCP< PHX::Evaluator<panzer::Traits> > je =
      Teuchos::rcp( new CLoadEvalutor<panzer::Traits::Jacobian, panzer::Traits>(cloads,globalIndexer) );
    phx_neumann_field_manager_->registerEvaluator<panzer::Traits::Jacobian>(je);
    phx_neumann_field_manager_->requireField<panzer::Traits::Jacobian>(*je->evaluatedFields()[0]);

    Teuchos::RCP< CLoadEvalutor<panzer::Traits::Tangent, panzer::Traits> > te =
      Teuchos::rcp( new CLoadEvalutor<panzer::Traits::Tangent, panzer::Traits>(cloads,globalIndexer) );
    phx_neumann_field_manager_->registerEvaluator<panzer::Traits::Tangent>(te);
    phx_neumann_field_manager_->requireField<panzer::Traits::Tangent>(*te->evaluatedFields()[0]);
  }

  panzer::Traits::SD setupData;
  phx_neumann_field_manager_->postRegistrationSetup(setupData);

  /*Teuchos::RCP<const panzer::GlobalIndexer> globalIndexer = lo_factory.getRangeGlobalIndexer();

          // Copy the physics block for side integrations
          Teuchos::RCP<panzer::PhysicsBlock> side_pb = volume_pb->copyWithCellData(side_cell_data);

          Teuchos::RCP<panzer::BCStrategy_TemplateManager<panzer::Traits> >
            bcstm = bc_factory.buildBCStrategy(*bc, side_pb->globalData());

          // Iterate over evaluation types
          int i=0;
          for (panzer::BCStrategy_TemplateManager<panzer::Traits>::iterator
                 bcs_type = bcstm->begin(); bcs_type != bcstm->end(); ++bcs_type,++i) {
            if (active_evaluation_types_[i]) {
              bcs_type->setDetailsIndex(block_id_index);
              side_pb->setDetailsIndex(block_id_index);
              bcs_type->setup(*side_pb, user_data);
              bcs_type->buildAndRegisterEvaluators(fm, *side_pb, cm_factory, closure_models, user_data);
              bcs_type->buildAndRegisterGatherAndOrientationEvaluators(fm, *side_pb, lo_factory, user_data);
              if ( ! physicsBlockScatterDisabled())
                bcs_type->buildAndRegisterScatterEvaluators(fm, *side_pb, lo_factory, user_data);
            }
          }

          gid_count += globalIndexer->getElementBlockGIDCount(element_block_id);
        }

        { // Use gid_count to set up the derivative information.
          std::vector<PHX::index_size_type> derivative_dimensions;
          derivative_dimensions.push_back(gid_count);
          fm.setKokkosExtendedDataTypeDimensions<panzer::Traits::Jacobian>(derivative_dimensions);

          #ifdef Panzer_BUILD_HESSIAN_SUPPORT
            fm.setKokkosExtendedDataTypeDimensions<panzer::Traits::Hessian>(derivative_dimensions);
          #endif

          derivative_dimensions[0] = 1;
          if (user_data.isType<int>("Tangent Dimension"))
            derivative_dimensions[0] = user_data.get<int>("Tangent Dimension");
          fm.setKokkosExtendedDataTypeDimensions<panzer::Traits::Tangent>(derivative_dimensions);
        }

        // Set up the field manager
        Traits::SD setupData;
        Teuchos::RCP<std::vector<panzer::Workset> > worksets = Teuchos::rcp(new std::vector<panzer::Workset>);
        worksets->push_back(wkst->second);
        setupData.worksets_ = worksets;
        setupData.orientations_ = getWorksetContainer()->getOrientations();

        Sacado::mpl::for_each_no_kokkos<panzer::Traits::EvalTypes>(PostRegistrationFunctor(active_evaluation_types_,fm,setupData));

      }
    } else {
      const std::string element_block_id = bc->elementBlockID();

      std::map<std::string,Teuchos::RCP<panzer::PhysicsBlock> >::const_iterator volume_pb_itr
	= physicsBlocks_map.find(element_block_id);

      TEUCHOS_TEST_FOR_EXCEPTION(volume_pb_itr==physicsBlocks_map.end(),std::logic_error,
				 "panzer::FMB::setupBCFieldManagers: Cannot find physics block corresponding to element block \"" << element_block_id << "\"");

      Teuchos::RCP<const panzer::PhysicsBlock> volume_pb = physicsBlocks_map.find(element_block_id)->second;
      Teuchos::RCP<const shards::CellTopology> volume_cell_topology = volume_pb->cellData().getCellTopology();

      // Build one FieldManager for each local side workset for each dirichlet bc
      std::map<unsigned,PHX::FieldManager<panzer::Traits> >& field_managers =
        bc_field_managers_[*bc];

      // Loop over local face indices and setup each field manager
      for (std::map<unsigned,panzer::Workset>::const_iterator wkst =
	     currentWkst->begin(); wkst != currentWkst->end();
	   ++wkst) {

        PHX::FieldManager<panzer::Traits>& fm = field_managers[wkst->first];

        // register evaluators from strategy
        const panzer::CellData side_cell_data(wkst->second.num_cells,
	                                      wkst->first,volume_cell_topology);

	// Copy the physics block for side integrations
	Teuchos::RCP<panzer::PhysicsBlock> side_pb = volume_pb->copyWithCellData(side_cell_data);

	Teuchos::RCP<panzer::BCStrategy_TemplateManager<panzer::Traits> > bcstm =
	  bc_factory.buildBCStrategy(*bc,side_pb->globalData());

	// Iterate over evaluation types
        int i=0;
	for (panzer::BCStrategy_TemplateManager<panzer::Traits>::iterator
	       bcs_type = bcstm->begin(); bcs_type != bcstm->end(); ++bcs_type,++i) {
          if (active_evaluation_types_[i]) {
            bcs_type->setup(*side_pb,user_data);
            bcs_type->buildAndRegisterEvaluators(fm,*side_pb,cm_factory,closure_models,user_data);
            bcs_type->buildAndRegisterGatherAndOrientationEvaluators(fm,*side_pb,lo_factory,user_data);
            if(!physicsBlockScatterDisabled())
              bcs_type->buildAndRegisterScatterEvaluators(fm,*side_pb,lo_factory,user_data);
          }
	}

	// Setup the fieldmanager
	Traits::SD setupData;
	Teuchos::RCP<std::vector<panzer::Workset> > worksets =
	  Teuchos::rcp(new(std::vector<panzer::Workset>));
	worksets->push_back(wkst->second);
	setupData.worksets_ = worksets;
        setupData.orientations_ = getWorksetContainer()->getOrientations();

	// setup derivative information
	setKokkosExtendedDataTypeDimensions(element_block_id,*globalIndexer,user_data,fm);

        Sacado::mpl::for_each_no_kokkos<panzer::Traits::EvalTypes>(PostRegistrationFunctor(active_evaluation_types_,fm,setupData));
      }
    }
  }*/
}

void FieldManagerBuilder::buildMaterials(const Teuchos::ParameterList& p)
{
      _materials.clear();
      for (auto it=p.begin(); it!= p.end(); ++it) {
         TEUCHOS_TEST_FOR_EXCEPTION( !(it->second.isList()), std::logic_error,
				"Error - All objects in the Material sublist must be unnamed sublists!" );
         Teuchos::ParameterList& sublist = Teuchos::getValue<Teuchos::ParameterList>(it->second);
    
         std::string matl_name = sublist.name(it);
         _materials[matl_name] = std::make_shared<Material>(sublist);
      }
}

}
